package statistics;

import java.util.Random;

/**
 * This class represents a mix of Geometric Distributions.
 */

public class MixGeometricDistribution extends DiscreteDistribution {

  private double p1;
  private double p2;
  private double q;

  /**
   * Constructs a mix of Geometric Distributions with parameters p1,q and p2,(1-q). You must also specify the
   * random seed that will be used to realize random variables with this distribution.
   * @param expect the expectation of the mix of Geometric Distributions.
   * @param a fit a mix of geometric distributions when a >= 1.
   */

  public MixGeometricDistribution(double expect, double a) {
    this(expect, a, new Random());
  }
  
  /**
   * Constructs a mix of Geometric Distributions with parameters p1,q and p2,(1-q). You must also specify the
   * random seed that will be used to realize random variables with this distribution.
   * @param expect the expectation of the mix of Geometric Distributions.
   * @param a fit a mix of geometric distributions when a >= 1.
   * @param random the randomgenerator .
   */

  public MixGeometricDistribution(double expect, double a, Random random) {
    this.random = random;
    p1 = expect * (1 + a + Math.sqrt(a*a - 1)) / (2 + expect*(1 + a + Math.sqrt(a*a - 1)));
    p2 = expect * (1 + a - Math.sqrt(a*a - 1)) / (2 + expect*(1 + a - Math.sqrt(a*a - 1)));
    q = 1.0/( 1 + a + Math.sqrt(a*a - 1) );
  }

  /**
   * Returns the expectation of a mix of Geometric distributions.
   */

  public double expectation() {
    return q*p1/(1-p1) + (1-q)*p2/(1-p2);
  }

  /**
   * Returns the variance of a mix of Geometric distributions.
   */

  public double variance() {
    return q*p1/((1-p1)*(1-p1)) + (1-q)*p2/((1-p2)*(1-p2));
  }

  /**
   * Returns P(X = k), X is a random variable with a Geometric distribution.
   * @param k
   */

  public double probabilityFunction(int k) {
    if (k < 0) return 0;
      else return q * (1.0-p1)*Math.pow(p1,k) + (1.0-q)*(1.0-p2)*Math.pow(p2,k);
  }

  /**
   * Returns a realization of a random variable with a mix of Geometric distributions.
   */

  public double nextRandom() {
    double u1 = random.nextDouble();
    double u = random.nextDouble();
    double outcome = 0;
    if(u1 < q) {
      outcome = Math.floor(Math.log(u)/Math.log(p1));
    }
    else {
      outcome = Math.floor(Math.log(u)/Math.log(p2));
    }
    return outcome;
  }
}
